import oneflow as torch
import oneflow.nn as nn


class LayerNorm(nn.LayerNorm):
    """Layer normalization module

    :param int nout: output dim size
    :param int dim: dimension to be normalized
    """

    def __init__(self, nout, dim=-1):
        super(LayerNorm, self).__init__(nout, eps=1e-12)
        self.dim = dim

    def forward(self, x):
        """Apply layer normalization

        :param torch.Tensor x: input tensor
        :return: layer normalized tensor
        :rtype torch.Tensor
        """
        if self.dim == -1:
            return super(LayerNorm, self).forward(x)
        return super(LayerNorm, self).forward(x.transpose(1, -1)).transpose(1, -1)


class LayerNorm2D(nn.Module):
    """Layer normalization for CNN outputs."""
    def __init__(self, dim, eps=1e-12):
        super(LayerNorm2D, self).__init__()
        self.norm = nn.LayerNorm(dim, eps=eps)

    def forward(self, xs):
        """Forward computation.
        Args:
            xs (FloatTensor): `[B, C, T, F]`
        Returns:
            xs (FloatTensor): `[B, C, T, F]`
        """
        B, C, T, F = xs.size()
        xs = xs.transpose(2, 1).contiguous().view(B, T, C * F)
        xs = self.norm(xs)
        xs = xs.view(B, T, C, F).transpose(2, 1)
        return xs